import os

import cv2
import numpy as np
from scipy.io import loadmat

from utils import moving_average

DebugMode = True
isCalibrate = False
if DebugMode:
    import matplotlib.pyplot as plt

if os.path.exists('mtx.mat'):
    isCalibrate = True
    data = loadmat('mtx.mat')
    mtx, newcameramtx, dist, roi = data['mtx'], data['new_mtx'], data['dist'], data['roi'][0]

calibrate_func = 'global y\n'
if os.path.exists('/home/pi/Desktop/func'):
    with open('/home/pi/Desktop/func', 'r') as f:
        calibrate_func += f.read()
else:
    calibrate_func += 'y = x * 1.1317 - 23.618'


def calibrate(img):
    if not isCalibrate:
        return img
    dst = cv2.undistort(img, mtx, dist, None, newcameramtx)
    x, y, w, h = roi
    return cv2.resize(dst[y:y + h, x:x + w], (img.shape[1], img.shape[0]))


def auto_brightness1(img: np.ndarray, out_range: tuple = (0, 255)) -> np.ndarray:
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    min_pixel_value, max_pixel_value = np.min(gray_img), np.max(gray_img)
    k = (out_range[1] - out_range[0]) / (max_pixel_value - min_pixel_value)
    b = out_range[0] - k * min_pixel_value
    img = np.clip((k * img + b), out_range[0], out_range[1])
    img = np.asarray(img, dtype=np.uint8)
    return img


def auto_brightness2(img: np.ndarray) -> np.ndarray:
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return cv2.equalizeHist(gray_img)


def get_water_level(pic: np.ndarray, rgb_mode=False, is_show=False) -> float:
    pic = calibrate(pic)
    if os.path.exists('mtx.mat'):
        mtx, newcameramtx, dist, roi = data['mtx'], data['new_mtx'], data['dist'], data['roi']
        dst = cv2.undistort(img, mtx, dist, None, newcameramtx)
    pic = pic if rgb_mode else pic[:, :, ::-1]  # convert color if not rgb mode

    # Get the ruler
    pic_hsv = cv2.cvtColor(pic, cv2.COLOR_RGB2HSV)
    plt.imshow(pic_hsv[:, :, 2])
    plt.title("v channel")
    plt.show()
    # th = cv2.adaptiveThreshold(pic_hsv[:, :, 0], 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 501, 5)
    # kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))  # 矩形结构
    # th = cv2.morphologyEx(th, cv2.MORPH_CLOSE, kernel)
    # ruler = pic[:, ~(th == 255).all(0)]
    ruler = pic[:, ((pic_hsv[:, :, 0] >= 20) & (pic_hsv[:, :, 0] <= 80) & (pic_hsv[:, :, 2] <= 240)).any(0), :]
    if DebugMode:
        plt.imshow(ruler)
        plt.title("Remove the background")
        _ = plt.show() if is_show else plt.savefig("debug01.png")

    # find the square block and E block
    ruler_v = cv2.cvtColor(ruler, cv2.COLOR_RGB2HSV)[:, :, 2]
    threshold_v = np.median(np.median(ruler)) * 0.7
    _, ruler_v_binary = cv2.threshold(ruler_v, threshold_v, maxval=255, type=cv2.THRESH_BINARY_INV)
    # find connected Components
    connected_num, connected_pic, stats, centers = cv2.connectedComponentsWithStats(ruler_v_binary, connectivity=8,
                                                                                    ltype=cv2.CV_32S)
    areas = stats[:, cv2.CC_STAT_AREA]
    areas[stats[:, cv2.CC_STAT_TOP] == 0] = 0
    max_area_idx = np.argmax(areas, axis=0)  # find the max E block
    e_block_top = stats[max_area_idx, cv2.CC_STAT_TOP]  # find the top of the E block
    e_block_bottom = stats[max_area_idx, cv2.CC_STAT_HEIGHT] + e_block_top
    if DebugMode:
        plt.imshow(ruler)
        plt.plot([0, ruler.shape[1] - 1], [e_block_top, e_block_top])
        plt.plot([0, ruler.shape[1] - 1], [e_block_bottom, e_block_bottom])
        plt.title("Find the Top and bottom of the E block")
        _ = plt.show() if is_show else plt.savefig("debug02.png")

    # find the water level line
    ruler_left_part = ruler[:, :int(ruler.shape[1] * 1), :]
    ruler_v_left = cv2.cvtColor(ruler_left_part, cv2.COLOR_RGB2HSV)[:, :, 2]
    threshold_v = np.median(np.median(ruler_left_part)) * 0.9
    _, ruler_v_left_binary = cv2.threshold(ruler_v_left, threshold_v, maxval=255, type=cv2.THRESH_BINARY_INV)
    avg_white = []
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))  # 矩形结构
    ruler_v_dilate = cv2.morphologyEx(ruler_v_left_binary, cv2.MORPH_DILATE, kernel)
    ruler_left_part[ruler_v_dilate != 0] = 127
    for row, white_part in zip(ruler_v_left, ruler_v_left_binary):
        avg_white.append(np.median(row[white_part == 0]))
    avg_white = np.array(avg_white)
    ruler_v[ruler_v_binary == 0] = 128
    avg_white = np.array(list(moving_average(avg_white, 7)))
    avg_white = avg_white[:-3] - avg_white[3:]
    avg_white[:e_block_bottom] = -100
    water_level_index = np.argmax(avg_white) + 3
    if DebugMode:
        fig, axs = plt.subplots(1, 1)
        axs.imshow(ruler)
        axs.plot(avg_white, np.arange(ruler_v.shape[0] - 9))
        axs.plot([0, ruler.shape[1] - 1], [water_level_index, water_level_index])
        axs.set_title("Find the water level Line")
        _ = plt.show() if is_show else plt.savefig("debug03.png")
    global y
    x = 210 - 50 / (e_block_bottom - e_block_top) * (water_level_index - e_block_bottom)
    exec(calibrate_func, globals(), locals())
    return y


if __name__ == '__main__':
    import os
    import time
    import argparse

    parser = argparse.ArgumentParser(description='Process some integers.')
    parser.add_argument("-d", "--debug", help="显示实时画面用于调整位置", action="store_true")
    parser.add_argument("-s", "--save", dest="save_to", required=False, type=str, default="",
                        help="保存一张图片,指定图片完整路径", metavar='FILE')
    parser.add_argument("-t", "--test", help="测试使用", required=False, type=str, default='')
    parser.add_argument("-n", "--notShow", help="not show the image", required=False, action='store_true')
    args = parser.parse_args()
    if args.save_to != "":
        from devs import Camera

        camera = Camera()
        img = camera.get_frame()
        cv2.imwrite(args.save_to, img)
        camera.stop()
    if args.debug:
        from devs import Camera

        camera = Camera()
        while True:
            img = camera.get_frame()
            # cv2.rectangle(img, pt1=(235, 218), pt2=(400, 360), color=[0, 0, 255], thickness=1)
            try:
                water_level = get_water_level(img)
            except Exception:
                water_level = -1
            cv2.putText(img, f"{water_level:.2f} cm", org=(20, 100), fontFace=cv2.FONT_HERSHEY_COMPLEX, fontScale=1, color=(0, 0, 255),
                        thickness=1)
            cv2.imshow(f"Press Q to quit, 按Q退出", img)
            key_code = cv2.waitKey(30)
            if (key_code & 0xFF == ord("q")) or (key_code & 0xFF == ord("Q")):
                break
    elif args.test != "":
        data_dir = args.test
        data_type = ["jpg", "png"]
        if not os.path.isdir(data_dir):
            imgs = [data_dir]
        else:
            imgs = [data_file for data_file in os.listdir(data_dir) if data_file.split(".")[-1] in data_type]
        print(f"Get imgs: {imgs}")
        print("START ===========================================")
        time_consumed_total = 0
        for img_name in imgs:
            since = time.time()
            img = cv2.imread(os.path.join(data_dir, img_name) if os.path.isdir(data_dir) else img_name)
            time_consumed = (time.time() - since)
            time_consumed_total += time_consumed
            water_level = get_water_level(img, is_show=True)
            print(water_level, img_name)
            # print(f"{img_name} water height: {water_level:.4f}, time consumed: {time_consumed:.4f} secs")
            # print(f"{img_name.split('-')[0]}.{img_name.split('-')[1]} | {water_level:4f}")
        print("END =============================================")
        avg_time = time_consumed_total / len(imgs)
        print(f"Average Time {avg_time:.4f} secs")
